from torchvision.models.resnet import BasicBlock as Res2DBasicBlock

import torch.nn as nn

'''
卷积神经网络的一些benchmark网络结构的标准块，包括半形块和保形块
'''

def half_hw_resnet_downsample_shortcut(in_channels, out_channels):
    return nn.Sequential(
        nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size=1,
            stride=2,
            bias=False), nn.BatchNorm2d(out_channels))

def half_hw_resblock(in_channels, out_channels):
    '''
    2D ResNet的半形块，输出的hw变为输入的1/2
    :param in_channels: 输入通道数
    :param out_channels: 输出通道数
    :return:
    '''
    return Res2DBasicBlock(in_channels, out_channels, stride=2,
               downsample=half_hw_resnet_downsample_shortcut(in_channels, out_channels))

def equal_hw_resblock(in_channels, out_channels=None):
    '''
    2D ResNet的保形块，块的输出尺寸与输出相同
    :param in_channels: 输入通道数
    :param out_channels: 输出通道数
    :return:
    '''
    if out_channels is None:
        out_channels = in_channels
    else:
        raise NotImplementedError
    return Res2DBasicBlock(in_channels, out_channels, stride=1,
               downsample=None)

def half_hw_convblock(in_channels, out_channels):
    '''
    2D 卷积半形块，输出的hw变为输入的1/2
    :param in_channels: 输入通道数
    :param out_channels: 输出通道数
    :return:
    '''
    return nn.Sequential(
    nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
    nn.BatchNorm2d(out_channels),
    nn.ReLU(inplace=True),
    nn.MaxPool2d(kernel_size=3, stride=2, padding=1))

def equal_hw_convblock(in_channels, out_channels):
    '''
    2D 卷积保形块，块的输出尺寸与输出相同
    :param in_channels: 输入通道数
    :param out_channels: 输出通道数
    :return:
    '''
    return nn.Sequential(
    nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
    nn.BatchNorm2d(out_channels),
    nn.ReLU(inplace=True))